import os
import cv2
import numpy as np
from typing import List, Dict, Any
from paddleocr import PaddleOCR
import tempfile


class OCREngine:
    """
    OCR文字识别引擎 基于PaddleOCR实现图片文字提取功能
    """
    
    def __init__(self):
        """
        初始化OCR引擎，使用指定的自定义模型路径
        """
        self.ocr = None
        self._initialize_ocr()
    
    def _initialize_ocr(self):
        """初始化OCR模型"""
        try:
            # 使用指定的自定义模型路径
            det_model_dir = 'D:/OCR_model/whl/det'
            rec_model_dir = 'D:/OCR_model/whl/rec'
            cls_model_dir = 'D:/OCR_model/whl/cls'
            
            # 检查模型路径是否存在
            if not os.path.exists(det_model_dir):
                raise Exception(f"检测模型路径不存在: {det_model_dir}")
            if not os.path.exists(rec_model_dir):
                raise Exception(f"识别模型路径不存在: {rec_model_dir}")
            if not os.path.exists(cls_model_dir):
                raise Exception(f"角度分类器模型路径不存在: {cls_model_dir}")
            
            # 使用自定义模型
            self.ocr = PaddleOCR(
                det_model_dir=det_model_dir,
                rec_model_dir=rec_model_dir,
                cls_model_dir=cls_model_dir,
                use_angle_cls=True,
                lang='ch',  # 支持中文
                show_log=False
            )
            
        except Exception as e:
            raise Exception(f"OCR模型初始化失败: {str(e)}")
    
    def recognize_from_file(self, image_path: str) -> Dict[str, Any]:
        """
        从图片文件识别文字
        
        Args:
            image_path: 图片文件路径
            
        Returns:
            包含识别结果的字典
        """
        if not os.path.exists(image_path):
            raise FileNotFoundError(f"图片文件不存在: {image_path}")
        
        try:
            # 执行OCR识别
            results = self.ocr.ocr(image_path, cls=True)
            
            # 解析结果
            parsed_results = self._parse_results(results)
            
            return {
                'success': True,
                'data': parsed_results,
                'image_path': image_path,
                'total_texts': len(parsed_results)
            }
            
        except Exception as e:
            return {
                'success': False,
                'error': str(e),
                'image_path': image_path
            }
    
    def recognize_from_bytes(self, image_bytes: bytes) -> Dict[str, Any]:
        """
        从图片字节数据识别文字
        
        Args:
            image_bytes: 图片字节数据
            
        Returns:
            包含识别结果的字典
        """
        try:
            # 将字节数据转换为numpy数组
            nparr = np.frombuffer(image_bytes, np.uint8)
            image = cv2.imdecode(nparr, cv2.IMREAD_COLOR)
            
            if image is None:
                raise ValueError("无法解码图片数据")
            
            # 创建临时文件
            with tempfile.NamedTemporaryFile(suffix='.jpg', delete=False) as tmp_file:
                tmp_path = tmp_file.name
                cv2.imwrite(tmp_path, image)
            
            try:
                # 使用临时文件进行识别
                result = self.recognize_from_file(tmp_path)
                return result
            finally:
                # 清理临时文件
                if os.path.exists(tmp_path):
                    os.unlink(tmp_path)
                    
        except Exception as e:
            return {
                'success': False,
                'error': str(e)
            }
    
    def _parse_results(self, ocr_results: List) -> List[Dict[str, Any]]:
        """
        解析OCR返回的原始结果
        
        Args:
            ocr_results: OCR原始结果
            
        Returns:
            解析后的结果列表
        """
        parsed_results = []
        
        if not ocr_results or not ocr_results[0]:
            return parsed_results
        
        for line in ocr_results[0]:
            if line and len(line) >= 2:
                # 提取文本和置信度
                text_info = line[-1]
                text = text_info[0]
                confidence = float(text_info[1])
                
                # 提取位置信息
                bbox = line[0]
                x1, y1 = int(bbox[0][0]), int(bbox[0][1])
                x2, y2 = int(bbox[2][0]), int(bbox[2][1])
                
                parsed_results.append({
                    'text': text,
                    'confidence': confidence,
                    'bbox': {
                        'x1': x1,
                        'y1': y1,
                        'x2': x2,
                        'y2': y2
                    },
                    'width': x2 - x1,
                    'height': y2 - y1
                })
        
        return parsed_results
    
    def extract_text_only(self, image_path: str) -> str:
        """
        仅提取文字内容，返回纯文本
        
        Args:
            image_path: 图片文件路径
            
        Returns:
            识别出的文字内容
        """
        result = self.recognize_from_file(image_path)
        
        if result['success']:
            texts = [item['text'] for item in result['data']]
            return '\n'.join(texts)
        else:
            return ""
    
    def batch_recognize(self, image_paths: List[str]) -> List[Dict[str, Any]]:
        """
        批量识别多张图片
        
        Args:
            image_paths: 图片路径列表
            
        Returns:
            批量识别结果
        """
        results = []
        for path in image_paths:
            result = self.recognize_from_file(path)
            results.append(result)
        return results
    
    def get_supported_languages(self) -> List[str]:
        """获取支持的语言列表"""
        return ['ch', 'en', 'fr', 'german', 'korean', 'japan']


# 全局实例
ocr_engine = OCREngine()